#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import logging
import re
import umap
import numpy as np
from sklearn.mixture import GaussianMixture
import trio

from graphrag.utils import (
    get_llm_cache,
    get_embed_cache,
    set_embed_cache,
    set_llm_cache,
    chat_limiter,
)
from rag.utils import truncate


# 定义一个用于递归抽象处理树形组织检索的类
class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
    # 初始化方法，接收最大聚类数、LLM模型、嵌入模型、提示模板、最大token数和阈值作为参数
    def __init__(
        self, max_cluster, llm_model, embd_model, prompt, max_token=512, threshold=0.1
    ):
        # 存储最大聚类数
        self._max_cluster = max_cluster
        # 存储大语言模型实例
        self._llm_model = llm_model
        # 存储嵌入模型实例
        self._embd_model = embd_model
        # 存储相似度阈值
        self._threshold = threshold
        # 存储提示模板
        self._prompt = prompt
        # 存储最大token数
        self._max_token = max_token

    # 异步聊天方法，接收系统提示、历史记录和生成配置
    async def _chat(self, system, history, gen_conf):
        # 从缓存中获取LLM响应，如果存在则直接返回
        response = get_llm_cache(self._llm_model.llm_name, system, history, gen_conf)
        if response:
            return response
        # 使用trio线程运行同步的LLM聊天方法
        response = await trio.to_thread.run_sync(
            lambda: self._llm_model.chat(system, history, gen_conf)
        )
        # 移除响应中的多余前缀（如果有）
        response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
        # 如果响应包含错误信息，则抛出异常
        if response.find("**ERROR**") >= 0:
            raise Exception(response)
        # 将响应结果存入缓存并返回
        set_llm_cache(self._llm_model.llm_name, system, response, history, gen_conf)
        return response

    # 异步嵌入编码方法，接收文本输入
    async def _embedding_encode(self, txt):
        # 从缓存中获取嵌入向量，如果存在则直接返回
        response = get_embed_cache(self._embd_model.llm_name, txt)
        if response is not None:
            return response
        # 使用trio线程运行同步的嵌入编码方法
        embds, _ = await trio.to_thread.run_sync(lambda: self._embd_model.encode([txt]))
        # 如果编码结果为空或格式不正确，则抛出异常
        if len(embds) < 1 or len(embds[0]) < 1:
            raise Exception("Embedding error: ")
        # 取出第一个文本的嵌入向量
        embds = embds[0]
        # 将嵌入向量存入缓存并返回
        set_embed_cache(self._embd_model.llm_name, txt, embds)
        return embds

    # 获取最优聚类数量，使用贝叶斯信息准则（BIC）评估不同聚类数的表现
    def _get_optimal_clusters(self, embeddings: np.ndarray, random_state: int):
        # 计算最大可能的聚类数，取预设最大值和样本数量中的较小者
        max_clusters = min(self._max_cluster, len(embeddings))
        # 定义要尝试的聚类数范围，从1到max_clusters-1
        n_clusters = np.arange(1, max_clusters)
        # 初始化一个列表用于存储每个聚类数对应的BIC值
        bics = []
        # 遍历每个聚类数，训练高斯混合模型并计算BIC值
        for n in n_clusters:
            # 初始化高斯混合模型（GMM）并设置随机种子
            gm = GaussianMixture(n_components=n, random_state=random_state)
            # 拟合嵌入向量数据
            gm.fit(embeddings)
            # 计算当前模型在数据上的BIC值，并添加到列表中
            bics.append(gm.bic(embeddings))
        # 找到最小BIC值对应的聚类数作为最优聚类数
        optimal_clusters = n_clusters[np.argmin(bics)]
        # 返回最优聚类数量
        return optimal_clusters

    # 异步调用方法，接收chunks（文本片段和嵌入）、随机状态和回调函数作为参数
    async def __call__(self, chunks, random_state, callback=None):
        # 如果chunks长度小于等于1，直接返回空列表
        if len(chunks) <= 1:
            return []
        # 过滤掉无效的chunk（s为空或a为空）
        chunks = [(s, a) for s, a in chunks if s and len(a) > 0]
        # 初始化层级信息，记录每层的起始和结束索引
        layers = [(0, len(chunks))]
        # 设置当前层的起始和结束位置
        start, end = 0, len(chunks)

        # 定义异步summarize函数，用于生成聚类摘要
        async def summarize(ck_idx: list[int]):
            # 声明chunks为外部作用域变量
            nonlocal chunks
            # 提取指定索引的文本内容
            texts = [chunks[i][0] for i in ck_idx]
            # 计算每个chunk可用的最大token数
            len_per_chunk = int(
                (self._llm_model.max_length - self._max_token) / len(texts)
            )
            # 将截断后的文本拼接成一个字符串
            cluster_content = "\n".join(
                [truncate(t, max(1, len_per_chunk)) for t in texts]
            )
            # 使用chat_limiter限制并发聊天请求
            async with chat_limiter:
                # 调用LLM生成摘要
                cnt = await self._chat(
                    "You're a helpful assistant.",
                    [
                        {
                            "role": "user",
                            "content": self._prompt.format(
                                cluster_content=cluster_content
                            ),
                        }
                    ],
                    {"temperature": 0.3, "max_tokens": self._max_token},
                )
            # 清理响应中的特定提示语句
            cnt = re.sub(
                "(······\n由于长度的原因，回答被截断了，要继续吗？|For the content length reason, it stopped, continue?)",
                "",
                cnt,
            )
            # 记录调试日志
            logging.debug(f"SUM: {cnt}")
            # 对生成的摘要进行编码
            embds = await self._embedding_encode(cnt)
            # 将摘要及其嵌入添加到chunks中
            chunks.append((cnt, embds))

        # 初始化标签列表
        labels = []
        # 循环处理每一层的内容
        while end - start > 1:
            # 获取当前层所有chunk的嵌入向量
            embeddings = [embd for _, embd in chunks[start:end]]
            # 如果只有两个chunk，则直接合并生成一个摘要
            if len(embeddings) == 2:
                # 调用summarize函数处理前两个chunk
                await summarize([start, start + 1])
                # 如果有回调函数，执行回调并传递进度信息
                if callback:
                    callback(
                        msg="Cluster one layer: {} -> {}".format(
                            end - start, len(chunks) - end
                        )
                    )
                # 添加标签
                labels.extend([0, 0])
                # 更新层级信息
                layers.append((end, len(chunks)))
                # 更新起始位置
                start = end
                # 更新结束位置
                end = len(chunks)
                # 继续下一轮循环
                continue

            # 计算邻居数量，用于UMAP降维
            n_neighbors = int((len(embeddings) - 1) ** 0.8)
            # 使用UMAP进行特征降维
            reduced_embeddings = umap.UMAP(
                n_neighbors=max(2, n_neighbors),
                n_components=min(12, len(embeddings) - 2),
                metric="cosine",
            ).fit_transform(embeddings)
            # 获取最优聚类数量
            n_clusters = self._get_optimal_clusters(reduced_embeddings, random_state)
            # 如果最优聚类数为1，则所有样本标记为同一类
            if n_clusters == 1:
                lbls = [0 for _ in range(len(reduced_embeddings))]
            else:
                # 初始化高斯混合模型并拟合数据
                gm = GaussianMixture(n_components=n_clusters, random_state=random_state)
                gm.fit(reduced_embeddings)
                # 预测每个样本属于各聚类的概率
                probs = gm.predict_proba(reduced_embeddings)
                # 根据阈值确定每个样本的聚类标签
                lbls = [np.where(prob > self._threshold)[0] for prob in probs]
                # 处理标签结果，确保是整数类型
                lbls = [lbl[0] if isinstance(lbl, np.ndarray) else lbl for lbl in lbls]

            # 创建一个新的协程组来并发执行summarize任务
            async with trio.open_nursery() as nursery:
                # 遍历每个聚类，并启动对应的summarize协程
                for c in range(n_clusters):
                    # 筛选属于当前聚类的chunk索引
                    ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c]
                    # 确保至少有一个chunk被分配给当前聚类
                    assert len(ck_idx) > 0
                    # 启动summarize协程
                    async with chat_limiter:
                        nursery.start_soon(summarize, ck_idx)

            # 确保新生成的chunk数量与聚类数一致
            assert len(chunks) - end == n_clusters, "{} vs. {}".format(
                len(chunks) - end, n_clusters
            )
            # 扩展标签列表
            labels.extend(lbls)
            # 更新层级信息
            layers.append((end, len(chunks)))
            # 如果有回调函数，执行回调并传递进度信息
            if callback:
                callback(
                    msg="Cluster one layer: {} -> {}".format(
                        end - start, len(chunks) - end
                    )
                )
            # 更新起始位置
            start = end
            # 更新结束位置
            end = len(chunks)

        # 返回处理后的chunks列表
        return chunks
